"""测试基于平面单应和平面方程已知的条件下求解相机位姿
已知：相机内参K、平面方程n、平面上的二维点y、像平面上的二维点x
"""

import numpy as np
import cv2
import argparse
from expriments import from_ellipse_args, ellipse_fit, ellipse_args
from expriments import EPnP, ED, ICP, Cone, EP, BA, PCCI, WP, CP
import tqdm
import os

method_list = [EPnP, ED, ICP, Cone, EP, BA, PCCI, WP, CP]


def arg_parse():
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument("--num", type=str, default=30, help="The number of points")
    arg_parser.add_argument("-f", type=float, default=25)
    arg_parser.add_argument("-d", type=float, default=5.5)
    arg_parser.add_argument("--reso", type=int, default=1024)
    arg_parser.add_argument("--noise", type=float, default=0)
    arg_parser.add_argument("--drift", type=float, default=0)
    arg_parser.add_argument("--height", type=int, default=4)
    arg_parser.add_argument(
        "--exp", type=int, default=100, help="The number of experiments"
    )
    arg_parser.add_argument("--recall", type=float, default=90)
    arg_parser.add_argument("--outlier", type=float, default=0)
    arg_parser.add_argument(
        "--use_init", action="store_true", help="Use pnp pose as initial"
    )
    arg_parser.add_argument("--debug", action="store_true", help="Debug mode")
    arg_parser.add_argument("--cache", type=str, default="test_crater_lib.npz")
    return arg_parser.parse_args()


if __name__ == "__main__":
    # 选用相同的随机种子，保证结果相对一致
    np.random.seed(114514)
    args = arg_parse()
    # 相机内参
    K = np.array(
        [
            [args.f / args.d * 1e3, 0, args.reso / 2],
            [0, args.f / args.d * 1e3, args.reso / 2],
            [0, 0, 1],
        ]
    )
    # 平面方程
    # n = np.array([0, 0, 1, 0])
    # 平面上任取三个点，构成生成子空间
    M = np.block([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
    if not args.cache or not os.path.exists(args.cache):
        # 平面上随机选取10000个圆, x,y,r
        elli_3d = []
        elli_2d = []
        C = []
        # 随机选取椭圆参数，生成仿真的陨石坑目录
        x = np.random.uniform(-4, 4, 2000) * 1e3
        y = np.random.uniform(-4, 4, 2000) * 1e3
        r = np.random.uniform(50, 200, 2000)
        for elli in zip(x, y, r, r):
            # 生成椭圆的二维点
            points_2d = from_ellipse_args(*elli, 0)
            points_2d = np.pad(
                points_2d, ((0, 1), (0, 0)), "constant", constant_values=1
            )
            # 二维点转三维点，以模拟相机成像过程
            points_3d = points_2d.T @ M
            C.append(ellipse_fit(points_2d[0], points_2d[1]))
            elli_3d.append(points_3d)
            elli_2d.append(points_2d)
        C = np.array(C)
        elli_3d = np.array(elli_3d)
        elli_2d = np.array(elli_2d)
        if args.cache:
            np.savez(args.cache, elli_3d=elli_3d, elli_2d=elli_2d, C=C)
    else:
        with np.load(args.cache) as data:
            elli_3d = data["elli_3d"]
            elli_2d = data["elli_2d"]
            C = data["C"]
    # 杂项设置
    test_num = 0
    progbar = tqdm.tqdm(total=args.exp, desc="Testing")
    failed = 0

    # 初始化位姿估计器
    method_kwargs = {"K": K, "use_pnp_initial": args.use_init}

    # 加载全部位姿估计器
    for it, method in enumerate(method_list):
        method_list[it] = method(**method_kwargs)
    # 初始化误差统计
    err = {m.__name__(): {} for m in method_list}
    while test_num < args.exp:
        # 生成像平面上的椭圆
        ## 相机随机位姿
        Rvec = np.random.randn(3, 1) * np.pi / 2
        R = cv2.Rodrigues(Rvec)[0]
        position = np.array([*np.random.uniform(-4, 4, 2), args.height]) * 1e3
        T = -R @ position
        P = K @ np.hstack((R, T[:, None]))
        # ## 生成像平面上的椭圆
        # H = P @ M.T
        # C_real = np.linalg.inv(H).T @ C @ np.linalg.inv(H)
        ## 在原始平面上采点，投影至像平面，再拟合为椭圆
        Qs = []
        Qs_prime = []
        obse_Qs_points = []
        for it, points_3d in enumerate(elli_3d):
            points_2d = cv2.projectPoints(
                points_3d[:, :3].T, Rvec, T, K.astype(np.float32), None
            )[0].squeeze()
            # 加入噪声和漂移
            points_2d += np.random.randn(*points_2d.shape) * args.noise + args.drift
            # 加入异常值点，即误检测的椭圆，同时删除一个正确的椭圆
            if np.random.rand() < args.outlier:
                r = np.random.uniform(10, 100, 2)
                theta = np.random.uniform(0, 2 * np.pi, 1)
                E = from_ellipse_args(
                    *np.block([np.random.uniform(0, 1024, 2), r, theta])
                )
                E = ellipse_fit(E[:, 0], E[:, 1])
            else:
                E = ellipse_fit(points_2d[:, 0], points_2d[:, 1])
                # C_real = cone_plane_cross(P.T @ E @ P, (0, 0, 1, 0))
            # 用当前椭圆的参数排除在视场以外的椭圆
            x, y, a, b, theta = ellipse_args(E)
            if x:
                if 0 <= x < args.reso and 0 <= y < args.reso:
                    Qs.append(E)
                    Qs_prime.append(C[it])
                    obse_Qs_points.append(from_ellipse_args(x, y, a, b, theta))
        # 考虑真实情况，同一幅图像中椭圆数量少于4个一定测不出位姿，大于20个不符合真实导航场景
        if len(Qs) < 4 or len(Qs) > 20:
            continue
        key = len(Qs)
        # 使用给定的求解器求解，返回求解时间
        for it, Estimator in enumerate(method_list):
            # 分别使用每种方法解算一次位姿，并保存至与其对应名为键的一级字典中，二级字典以当前椭圆数量为键
            time, success, rot, tvec = Estimator(Qs_prime, Qs, obse_Qs_points)
            name = Estimator.__name__()
            if key not in err[name].keys():
                err[name][key] = {"fail": 0, "err": [], "total": 0, "time": []}
            now_err_dict = err[name][key]
            now_err_dict["total"] += 1
            if not success:
                now_err_dict["fail"] += 1
                continue
            # 消除二义性，因为相机的世界坐标一定位于平面上方，且世界坐标系符合右手系
            esti_position = -(rot.T @ tvec).squeeze()
            if esti_position[-1] < 0:
                rot[:, :2] = -rot[:, :2]
                esti_position = (rot.T @ tvec).squeeze()
            # rvec = cv2.Rodrigues(rot)[0]
            now_err_dict["time"].append(time)

            err_dist = np.power(position - esti_position, 2).sum()
            err_rot = (
                np.power(cv2.Rodrigues(rot.T @ R)[0], 2).sum() * (180 / np.pi) ** 2
            )
            now_err_dict["err"].append((err_dist, err_rot))
        test_num += 1
        progbar.update(1)

    progbar.close()
    # 写出至文件流中，如果是调试模式，则将当前所有数据都写入至文件流中
    for Estimator in method_list:
        total_dist = []
        total_rot = []
        times = []
        total_failed = 0
        name = Estimator.__name__()
        if args.debug:
            with open(
                f"exp-{int(args.noise)}-{int(args.drift)}-{args.outlier}-{name}.csv",
                "w",
            ) as w:
                now_err_dict = err[name]
                print("num of craters,success/%,dist,rot,run time/ms", file=w)
                for k, v in now_err_dict.items():
                    success = 1 - v["fail"] / v["total"]
                    total_failed += v["fail"]
                    if len(v["err"]) == 0:
                        continue
                    err_list = np.array(v["err"])
                    dist = np.sqrt(err_list[:, 0].mean())
                    rot = np.sqrt(err_list[:, 1].mean())
                    time = np.mean(v["time"])
                    print(f"{k},{success},{dist},{rot},{time*1e3}", file=w, flush=True)
                    total_dist.append(dist)
                    total_rot.append(rot)
                    times.append(time)
                # 覆盖当前字典以节省内存
                err[name] = {
                    "dist": np.array(total_dist),
                    "rot": np.array(total_rot),
                    "time": np.array(times),
                    "failed": total_failed,
                }
        else:
            now_err_dict = err[name]
            for k, v in now_err_dict.items():
                success = 1 - v["fail"] / v["total"]
                total_failed += v["fail"]
                if len(v["err"]) == 0:
                    continue
                err_list = np.array(v["err"])
                dist = np.sqrt(err_list[:, 0].mean())
                rot = np.sqrt(err_list[:, 1].mean())
                time = np.mean(v["time"])
                total_dist.append(dist)
                total_rot.append(rot)
                times.append(time)
            # 覆盖当前字典以节省内存
            err[name] = {
                "dist": np.array(total_dist),
                "rot": np.array(total_rot),
                "time": np.array(times),
                "failed": total_failed,
            }

    print(f"R@{args.recall/100}\t\tsuccess\tdist\trot\trun time/ms")
    # with open("result.csv", "a") as w:
    #     print(
    #         f"{int(args.noise)}/{int(args.drift)}/{args.outlier}/{args.height},success,dist,rot,runtime/ms",
    #         file=w,
    #     )
    for Estimator in method_list:
        name = Estimator.__name__()
        total_dist = err[name]["dist"]
        total_rot = err[name]["rot"]
        times = err[name]["time"]
        total_failed = err[name]["failed"]
        # 计算90%分位数的误差，排除奇异点
        total_dist = total_dist[total_dist < np.percentile(total_dist, args.recall)]
        dist = np.sqrt(np.power(total_dist, 2).mean())

        total_rot = total_rot[total_rot < np.percentile(total_rot, args.recall)]
        rot = np.sqrt(np.power(total_rot, 2).mean())
        print(
            f"{name}\t\t{1-total_failed/args.exp:.3f}\t{dist:.3f}\t{rot:.3f}\t{1000*times.mean():.3f}"
        )
        # print(
        #     f"{name},{1-total_failed/args.exp:.3f},{dist:.3f},{rot:.3f},{1000*times.mean():.3f}",
        #     file=w,
        # )
    print(
        f"R@{args.recall/100}\tnoise\tdrift\toutlier\theight\texpriment times\tsuccess\tdist\trot\trun time/ms"
    )
    for Estimator in method_list:
        # 同一种方法输出到同一文件中
        name = Estimator.__name__()
        with open(f"results/{name}.csv", "a") as w:
            w.write(
                f"{args.noise},{args.drift},{args.outlier},{args.height},{args.exp},"
            )
            total_dist = err[name]["dist"]
            total_rot = err[name]["rot"]
            times = err[name]["time"]
            total_failed = err[name]["failed"]
            # 计算90%分位数的误差，排除奇异点
            total_dist = total_dist[total_dist < np.percentile(total_dist, args.recall)]
            dist = np.sqrt(np.power(total_dist, 2).mean())

            total_rot = total_rot[total_rot < np.percentile(total_rot, args.recall)]
            rot = np.sqrt(np.power(total_rot, 2).mean())
            w.write(
                f"{1-total_failed/args.exp:.3f},{dist:.3f},{rot:.3f},{1000*times.mean():.3f}\n"
            )
